# -*- coding: utf-8 -*-
import pendulum
import random
from datetime import datetime, timedelta
import math

cst = pendulum.timezone('Asia/Shanghai')


def random_choice(seq):
    return random.choice(seq)


def cst_now() -> datetime:
    return cst.convert(datetime.now())


def date_add(execution_date: datetime, days: int) -> datetime:
    return execution_date + timedelta(days=days)

def hour_add(execution_date: datetime, hours: int) -> datetime:
    return execution_date + timedelta(hours=hours)

def week_add(execution_date: datetime, weeks: int) -> datetime:
    return execution_date + timedelta(weeks=weeks)

def to_cst(execution_date: datetime) -> datetime:
    return cst.convert(execution_date)


def cst_format(execution_date: datetime, fmt: str) -> str:
    return to_cst(execution_date).strftime(fmt)


def cst_ds(execution_date: datetime) -> str:
    return cst_format(execution_date, '%Y-%m-%d')

def cst_hour(execution_date: datetime) -> str:
    return cst_format(execution_date, '%Y-%m-%d-%H')

def cst_hour_rt(execution_date: datetime) -> str:
    return cst_format(execution_date, '%Y-%m-%d--%H')

def route_cst_ds(execution_date: datetime) -> str:
    test_date = datetime(year=execution_date.year, month=3, day=execution_date.day)
    return cst_format(test_date, '%Y-%m-%d')

def cst_ds_sfm(execution_date: datetime) -> str:
    return cst_format(execution_date, '%Y-%m-%d %H:%M:%S')

def cst_hour_ds(execution_date: datetime) -> str:
    return cst_format(execution_date, '%Y-%m-%d %H')


def cst_ds_nodash(execution_date: datetime) -> str:
    return cst_format(execution_date, '%Y%m%d')


def cst_weekday(execution_date: datetime) -> str:
    return cst_format(execution_date, '%w')

def cst_week(execution_date: datetime) -> str:
    return cst_format(execution_date,'%Y-%W')

def cstweek(execution_date: datetime) -> str:
    return cst_format(execution_date,'%Y%W')

def cst_monthday(execution_date: datetime) -> str:
    return cst_format(execution_date, '%d')


def upper(s):
    return str.upper(s)


def lower(s):
    return str.lower(s)

def cst_hs(execution_date: datetime) -> str:
    return cst_format(execution_date, '%Y-%m-%d-%H')
#月函数
def cst_month_dash(execution_date: datetime) -> str:
    return cst_format(execution_date, '%Y-%m')

def cst_month_nodash(execution_date: datetime) -> str:
    return cst_format(execution_date, '%Y%m')

def add_month(execution_date: datetime, num: int) -> datetime:
    months = ['01', '02', '03', '04', '05', '06', '07', '08', '09', '10', '11', '12']
    datamonth=cst_ds_nodash(execution_date)
    year_month = int(datamonth[0:6])
    num = int(num)
    year = year_month // 100
    new_list = []
    s = math.ceil(abs(num) / 12)
    for i in range(int(-s), s + 1):
        new_list += [str(year + i) + x for x in months]
    new_list = [int(x) for x in new_list]
    yearMonth=str(new_list[new_list.index(year_month) + num])
    result = yearMonth[0:4]+"-"+yearMonth[4:6]
    month = datetime.strptime(result,"%Y-%m")
    return month

#周  execution_date  "2022-12-31 01:01:01"  2021-52
def year_week(execution_date: datetime)-> str:
    y=execution_date.year
    m=execution_date.month
    d=execution_date.day
    yw=""
    week=datetime(int(y), int(m), int(d)).isocalendar()[1]    #周
    year_of_week=str(datetime(int(y), int(m), int(d)).isocalendar()[0]) #年
    #周是个位数，则在周前加0
    week_str=str(week)
    if(week<10):
        week_str="0"+week_str
        yw=year_of_week+"-"+week_str
    yw=year_of_week+"-"+week_str
    return yw

def replace_str(str,old,new)->str:
    return str.replace(old,new)

macros_dict = {'random_choice': random_choice,
               'cst_now'      : cst_now,
               'dt_add'       : date_add,  # alias
               'date_add'     : date_add,
               'cst_format'   : cst_format,
               'cst_ds'       : cst_ds,
               'cst_hs'       : cst_hs,            #新增2022-05-25
               'cst_ds_nodash': cst_ds_nodash,
               'to_cst'       : to_cst,
               'cst'          : to_cst,  # alias
               'upper'        : upper,
               'lower'        : lower,
               'cst_month_dash':cst_month_dash,     #新增2022-05-25
               'cst_month_nodash':cst_month_nodash, #新增2022-05-25
               'route_cst_ds' : route_cst_ds,
               'hour_add'     : hour_add,
               'cst_hour'     : cst_hour,
               'cst_hour_rt'  : cst_hour_rt,
               'add_month'    :add_month,           #新增2022-05-25
               'year_week'    : year_week,          #新增2022-05-25
               'replace_str'  : replace_str,         #新增2022-05-25
               'cst_ds_sfm'   : cst_ds_sfm,
               'cst_hour_ds'  : cst_hour_ds     #新增2022-06-16
               }
